#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
ESP32服务端模拟器
实现混合加密的TCP服务端
"""

import socket
import threading
import json
import base64
import os
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.fernet import Fernet
import time

class ESP32Server:
    def __init__(self, host='localhost', port=8888):
        self.host = host
        self.port = port
        self.server_socket = None
        self.private_key = None
        self.public_key = None
        self.clients = {}  # 存储客户端连接和对应的AES密钥
        self.running = False
        
        # 加载RSA密钥
        self.load_rsa_keys()
    
    def load_rsa_keys(self):
        """
        加载RSA密钥对
        """
        try:
            # 加载私钥
            with open('keys/server_private_key.pem', 'rb') as f:
                self.private_key = serialization.load_pem_private_key(
                    f.read(),
                    password=None
                )
            
            # 加载公钥
            with open('keys/server_public_key.pem', 'rb') as f:
                self.public_key = serialization.load_pem_public_key(f.read())
            
            print("RSA密钥加载成功")
        except FileNotFoundError:
            print("错误：找不到密钥文件，请先运行 generate_keys.py 生成密钥")
            exit(1)
    
    def start_server(self):
        """
        启动服务器
        """
        try:
            self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.server_socket.bind((self.host, self.port))
            self.server_socket.listen(5)
            self.running = True
            
            print(f"ESP32服务器启动成功，监听 {self.host}:{self.port}")
            print("等待客户端连接...")
            
            while self.running:
                try:
                    client_socket, client_address = self.server_socket.accept()
                    print(f"新客户端连接：{client_address}")
                    
                    # 为每个客户端创建处理线程
                    client_thread = threading.Thread(
                        target=self.handle_client,
                        args=(client_socket, client_address)
                    )
                    client_thread.daemon = True
                    client_thread.start()
                    
                except socket.error as e:
                    if self.running:
                        print(f"接受连接时出错：{e}")
                    break
                    
        except Exception as e:
            print(f"服务器启动失败：{e}")
        finally:
            self.stop_server()
    
    def handle_client(self, client_socket, client_address):
        """
        处理客户端连接
        """
        try:
            # 1. 发送公钥给客户端
            public_key_pem = self.public_key.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo
            )
            
            response = {
                'type': 'public_key',
                'data': base64.b64encode(public_key_pem).decode('utf-8')
            }
            self.send_message(client_socket, response)
            print(f"已向 {client_address} 发送公钥")
            
            # 2. 接收客户端加密的AES密钥
            encrypted_aes_key_msg = self.receive_message(client_socket)
            if encrypted_aes_key_msg and encrypted_aes_key_msg['type'] == 'encrypted_aes_key':
                encrypted_aes_key = base64.b64decode(encrypted_aes_key_msg['data'])
                
                # 使用RSA私钥解密AES密钥
                aes_key = self.private_key.decrypt(
                    encrypted_aes_key,
                    padding.OAEP(
                        mgf=padding.MGF1(algorithm=hashes.SHA256()),
                        algorithm=hashes.SHA256(),
                        label=None
                    )
                )
                
                # 存储客户端的AES密钥
                fernet = Fernet(aes_key)
                self.clients[client_address] = {
                    'socket': client_socket,
                    'fernet': fernet
                }
                
                print(f"已建立与 {client_address} 的加密通道")
                
                # 发送确认消息
                confirm_msg = {'type': 'handshake_complete', 'message': '加密通道建立成功'}
                self.send_encrypted_message(client_socket, fernet, confirm_msg)
                
                # 3. 处理后续的加密通信
                self.handle_encrypted_communication(client_socket, client_address, fernet)
            
        except Exception as e:
            print(f"处理客户端 {client_address} 时出错：{e}")
        finally:
            # 清理客户端连接
            if client_address in self.clients:
                del self.clients[client_address]
            client_socket.close()
            print(f"客户端 {client_address} 断开连接")
    
    def handle_encrypted_communication(self, client_socket, client_address, fernet):
        """
        处理加密通信
        """
        while self.running:
            try:
                # 接收加密消息
                encrypted_msg = self.receive_message(client_socket)
                if not encrypted_msg:
                    break
                
                if encrypted_msg['type'] == 'encrypted_data':
                    # 解密消息
                    encrypted_data = base64.b64decode(encrypted_msg['data'])
                    decrypted_data = fernet.decrypt(encrypted_data)
                    message = json.loads(decrypted_data.decode('utf-8'))
                    
                    print(f"收到来自 {client_address} 的消息：{message}")
                    
                    # 处理不同类型的消息
                    response = self.process_message(message)
                    
                    # 发送加密响应
                    self.send_encrypted_message(client_socket, fernet, response)
                    
            except Exception as e:
                print(f"处理加密通信时出错：{e}")
                break
    
    def process_message(self, message):
        """
        处理接收到的消息
        """
        msg_type = message.get('type', 'unknown')
        
        if msg_type == 'ping':
            return {
                'type': 'pong',
                'message': 'ESP32服务器响应',
                'timestamp': time.time()
            }
        elif msg_type == 'sensor_request':
            # 模拟传感器数据
            return {
                'type': 'sensor_data',
                'temperature': 25.6,
                'humidity': 60.2,
                'timestamp': time.time()
            }
        elif msg_type == 'control':
            action = message.get('action', '')
            return {
                'type': 'control_response',
                'action': action,
                'status': 'success',
                'message': f'已执行操作：{action}'
            }
        else:
            return {
                'type': 'error',
                'message': f'未知消息类型：{msg_type}'
            }
    
    def send_message(self, client_socket, message):
        """
        发送JSON消息
        """
        try:
            data = json.dumps(message).encode('utf-8')
            length = len(data)
            client_socket.sendall(length.to_bytes(4, byteorder='big'))
            client_socket.sendall(data)
        except Exception as e:
            print(f"发送消息失败：{e}")
    
    def receive_message(self, client_socket):
        """
        接收JSON消息
        """
        try:
            # 先接收消息长度
            length_bytes = client_socket.recv(4)
            if not length_bytes:
                return None
            
            length = int.from_bytes(length_bytes, byteorder='big')
            
            # 接收消息数据
            data = b''
            while len(data) < length:
                chunk = client_socket.recv(length - len(data))
                if not chunk:
                    return None
                data += chunk
            
            return json.loads(data.decode('utf-8'))
        except Exception as e:
            print(f"接收消息失败：{e}")
            return None
    
    def send_encrypted_message(self, client_socket, fernet, message):
        """
        发送加密消息
        """
        try:
            # 加密消息
            message_json = json.dumps(message).encode('utf-8')
            encrypted_data = fernet.encrypt(message_json)
            
            # 发送加密消息
            encrypted_msg = {
                'type': 'encrypted_data',
                'data': base64.b64encode(encrypted_data).decode('utf-8')
            }
            self.send_message(client_socket, encrypted_msg)
        except Exception as e:
            print(f"发送加密消息失败：{e}")
    
    def stop_server(self):
        """
        停止服务器
        """
        self.running = False
        if self.server_socket:
            self.server_socket.close()
        print("服务器已停止")

def main():
    print("=== ESP32服务端模拟器 ===")
    
    # 检查密钥文件是否存在
    if not os.path.exists('keys/server_private_key.pem'):
        print("错误：找不到密钥文件，请先运行 generate_keys.py 生成密钥")
        return
    
    server = ESP32Server()
    
    try:
        server.start_server()
    except KeyboardInterrupt:
        print("\n收到中断信号，正在关闭服务器...")
        server.stop_server()

if __name__ == "__main__":
    main()